Scripts and a skill to do per-layer benchmark using flashinfer#1980
Scripts and a skill to do per-layer benchmark using flashinfer#1980sychen52 wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdds an operator-facing kernel benchmarking skill, meta-tensor inspection for dense, Mamba, and MoE layouts, FlashInfer execution with quantization timing and error reporting, comprehensive tests, and Bandit exclusions. ChangesModel Kernel Benchmarking
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Operator
participant benchmark_model.py
participant Transformers
participant benchmark_via_builtin.py
participant flashinfer_benchmark.py
Operator->>benchmark_model.py: Select model, TP/EP, and M values
benchmark_model.py->>Transformers: Construct meta-tensor model
Transformers-->>benchmark_model.py: Return module and config shapes
benchmark_model.py->>benchmark_via_builtin.py: Pass derived GEMM and MoE arguments
benchmark_via_builtin.py->>flashinfer_benchmark.py: Run generated testlist cases
flashinfer_benchmark.py-->>benchmark_via_builtin.py: Return timings and errors
benchmark_via_builtin.py-->>Operator: Write combined_results.csv
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py:
- Around line 274-286: The argument parsers in benchmark_model.py (around parser
construction before _load_meta_model) and benchmark_via_builtin.py (around its
main parser) must validate all numeric benchmark inputs as positive integers.
Add shared positive-integer validation for --tp and --ep in benchmark_model.py,
and for M/N/K, MoE shape, and iteration arguments in benchmark_via_builtin.py,
so invalid non-positive values are rejected during parsing before model
inspection, shape generation, execution, or CUDA allocation.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 399-426: Update the CSV-writing flow in the results writer to emit
the separator rows required by the sectioned-table contract around both the GEMM
and MoE sections. Preserve the existing headers, data rows, and conditional
section ordering while restoring separators even when both sections are present.
- Around line 329-337: Update the FP8 branch in _quant_times() so that when
vllm_ops is unavailable it records an explicit warning or error result for the
affected quantization case instead of continuing silently. Ensure
_write_results() preserves and emits that fallback cell as the corresponding
_with_quant row, while leaving available vLLM and NVFP4 execution unchanged.
- Around line 525-537: Update the result-loading block around _run_driver and
builtin_csv so a missing builtin_results.csv is treated as an empty row set
instead of raising FileNotFoundError. Preserve normal CSV parsing when the file
exists, then continue completed_tags and per-case error aggregation so
combined_results.csv is still produced with missing-case reasons.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8a542345-e207-4bbb-9999-64efe1192a93
📒 Files selected for processing (7)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/agents/openai.yaml.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/evals.json.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py
4d9cabf to
9a63807
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1980 +/- ##
==========================================
- Coverage 77.87% 74.20% -3.68%
==========================================
Files 522 525 +3
Lines 58452 62383 +3931
==========================================
+ Hits 45522 46289 +767
- Misses 12930 16094 +3164
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 350-374: Update the per-spec quantization timing loop in
_quant_times to catch CUDA allocation, CUPTI, runner, and kernel timing
exceptions without aborting the overall benchmark. For each failed (kind, m, k)
specification, store an explicit error result using the existing error-value
convention like _FP8_QUANT_UNAVAILABLE, emit an appropriate warning if needed,
and continue so completed driver results can still be written to
combined_results.csv.
In `@pyproject.toml`:
- Around line 291-297: Remove the .agents/skills/benchmark-model-kernels/ entry
from Bandit’s exclude_dirs configuration so executable benchmark scripts remain
covered. Add a narrow suppression only at the specific reviewed false-positive
call site within those scripts, without excluding the broader skill directory.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: db8b82d3-7d5e-4dbd-8b91-1a668cabacc9
📒 Files selected for processing (8)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/agents/openai.yaml.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/evals.json.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.pypyproject.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- .agents/skills/benchmark-model-kernels/agents/openai.yaml
- .agents/skills/benchmark-model-kernels/tests/evals.json
- .agents/skills/benchmark-model-kernels/SKILL.md
| exclude_dirs = [ | ||
| ".agents/skills/benchmark-model-kernels/", | ||
| ".github/", | ||
| "examples/", | ||
| "noxfile.py", | ||
| "tests/", | ||
| ] |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Keep the executable benchmark scripts in Bandit coverage.
This excludes the entire skill, including its production scripts. Remove Line 292 and narrowly suppress any reviewed false positive at its call site instead.
Proposed fix
exclude_dirs = [
- ".agents/skills/benchmark-model-kernels/",
".github/",
"examples/",
"noxfile.py",
"tests/",
]As per coding guidelines, “Make simple, surgical changes; avoid speculative refactors, broad rewrites, and unrelated cleanup.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| exclude_dirs = [ | |
| ".agents/skills/benchmark-model-kernels/", | |
| ".github/", | |
| "examples/", | |
| "noxfile.py", | |
| "tests/", | |
| ] | |
| exclude_dirs = [ | |
| ".github/", | |
| "examples/", | |
| "noxfile.py", | |
| "tests/", | |
| ] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pyproject.toml` around lines 291 - 297, Remove the
.agents/skills/benchmark-model-kernels/ entry from Bandit’s exclude_dirs
configuration so executable benchmark scripts remain covered. Add a narrow
suppression only at the specific reviewed false-positive call site within those
scripts, without excluding the broader skill directory.
Source: Coding guidelines
There was a problem hiding this comment.
@kevalmorabia97 @shengliangxu , I added a python script that will call the benchmark utils in flashinfer to run benchmark, so it will use subprocess to call it. Can I add this script to the exclude directories of bandit?
There was a problem hiding this comment.
why? does bandit complain about subprocess?
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
@yeyu-nvidia can you just add a # nosec: <CODE> next to all lines where bandit complains instead of skipping the folder? If there are too many then we can just keep it here
There was a problem hiding this comment.
@shengliangxu yes.
@kevalmorabia97 I did that initially, but CodeRabbit's pre-merge check give me an error. I guess you mean I add # nosec and disregard the error and you will approval it in the PR, right?
with different backends. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
9a63807 to
2cd8630
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 33-36: Update the import and subprocess setup around argument
parsing, `run_benchmark`, and the child execution flow so `--flashinfer_repo`
controls both `flashinfer` and `flashinfer.testing` resolution. Add the selected
checkout root to `sys.path` before importing those modules, including when the
child runs from `benchmarks_dir`, and ensure imports cannot silently use an
unrelated installed FlashInfer.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9c1b8ca4-8092-4081-b771-83772f0e2180
📒 Files selected for processing (6)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.pypyproject.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- pyproject.toml
- .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py
- .agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py
sychen52
changed the title
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-4-8) — DM the bot to share feedback.
New self-contained skill (benchmark-model-kernels) with two scripts that (1) derive per-rank GEMM/MoE benchmark shapes from a HF config on meta tensors and (2) drive FlashInfer's built-in benchmark for BF16/FP8/NVFP4. +2301/-1, 3 non-test files + 2 test files + a bandit-exclude tweak.
Design review (complexity gate fired): The problem — a per-layer microbenchmark harness that derives per-rank shapes from model configs without loading weights — has no existing in-repo equivalent (searched; other skills are task-specific and use tests.json eval metadata, not benchmarking). This wraps FlashInfer's flashinfer_benchmark.py rather than reimplementing it; the value-add is the config→shape derivation. OmegaConf/Hydra/Jinja aren't relevant. The pytest-in-skill layout is already an established convention (pyproject has a .agents/skills/*/tests/test_*.py per-file-ignore). Design looks reasonable and the PR body gives concrete usage, though the rationale is brief.
Licensing: clean — new files carry the canonical LICENSE_HEADER NVIDIA header verbatim; no LICENSE/vendoring changes. The bandit exclude_dirs addition is narrowly scoped to the new skill (justified by the intentional subprocess use, which is invoked without a shell and gates trust_remote_code behind an explicit flag).
Correctness: I traced the shape-derivation logic (fused QKV incl. GQA replication, gate/up, Mamba-2, ModuleList vs stacked/direct experts, EP-vs-TP intermediate sharding, routing derivation, expert audit, override rejection) and it holds up; the unit tests are genuinely meaningful and internally consistent across Llama/Mixtral/GptOss/NemotronH plus edge cases.
Why nudge (not approve): (1) sizable new subsystem worth an owner's architectural eyes; (2) coverage gap — benchmark_via_builtin.py's GPU/FlashInfer runtime paths (_run_driver real invocation, _quant_times/_nvfp4_runner/_fp8_runner CUDA timing) are mocked out and never exercised in CI, consistent with the repo's known GPU-only coverage risk; (3) CHANGELOG left unmarked — acceptable for an internal agent skill (no shipped library API change), but worth confirming. No prompt-injection concerns in the untrusted content (SKILL.md's "do not override" prose targets the skill's agent user, not the reviewer).
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Posting two inline comments from my review: one on CI test discovery, one on TP×EP handling for the MoE shape.
| pytest.importorskip("torch") | ||
| pytest.importorskip("accelerate") | ||
| transformers = pytest.importorskip("transformers") |
There was a problem hiding this comment.
These tests never run in CI as written. The skills job in .github/workflows/unit_tests.yml installs only pytest (it's documented as "stdlib-only and hermetic"), so these importorskip calls on torch/accelerate/transformers skip this entire file — and pytest.importorskip("flashinfer") skips all of test_benchmark_via_builtin.py the same way. CI stays green whether or not the code works, so the ~870 lines of tests only ever run on a dev machine.
Suggestions:
- Route this file into the main
linuxunit lane (it needs only torch/accelerate/transformers — all standard unit-test deps, no GPU). - In
benchmark_via_builtin.py, deferimport flashinferinto the functions that need it (same pattern as the existingvllmtry/except). The majority of that test file (_parser,_combine,_write_results,_parse_driver_errors, case generation) is pure CPU logic that could then run without the GPU stack.
There was a problem hiding this comment.
import moved. However, route test to main unit lane involve workflow change.
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
5 participants
add scripts and a skill to use flashinfer to do layerwise benchmark with different backends.
What does this PR do?
Type of change: ? new skill and scripts
Usage
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ openai/gpt-oss-120b \ --tp 4 --ep 4 --ms 8 16 \ --flashinfer_repo $HOME/flashinfer \ --workdir /tmp/gpt-oss-benchmarkor
python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \ --flashinfer_repo $HOME/flashinfer \ --ms 8 16 \ --nks 1280,2880 2880,1024 \ --moe_hidden_size 2880 \ --moe_intermediate_size 2880 \ --moe_num_experts 32 \ --moe_top_k 4 \ --workdir /tmp/gpt-oss-via-builtinTesting
run locally
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit
benchmark-model-kernelsskill with a single-decision workflow for planning and running BF16/FP8/NVFP4 GEMM and fused-MoE microbenchmarks.combined_results.csvoutput with quantization and error details.